import sys
import json
import cv2
import math
import numpy as np
import torch as pt
from pathlib import Path
from nex_network import NexMsi, basis_order, basis_out, pos_level, sigmoid_offset

DEBUG = False

# export msi images, range: [0, 1]
def export_msi_images(msi_rgb, output_folder):
  msi_rgb_sig = pt.sigmoid(msi_rgb)
  np.save(output_folder / 'msi_rgb', msi_rgb_sig.cpu().numpy())
  if DEBUG:
    for i, rgb_sig in enumerate(msi_rgb_sig):
      im = (rgb_sig[0] * 255 + 0.5).to(pt.uint8).cpu().numpy()
      im = cv2.cvtColor(im, cv2.COLOR_RGB2BGR)
      cv2.imwrite(str(output_folder / f'msi_rgb_{i:02d}.png'), im)

# export nex bases, range: [-1, 1]
# use pixel on 'The Equi-Angular Cubemap' as directions
# cubemap resolution: [resolution, resolution, 6]
# input opencv coordinate
# exported cubemaps are converted to opengl cubemap coordinate
def export_nex_bases(model_specular, output_folder, resolution):
  def xyzs_to_theta_phi(xyzs): # [..., 3]
    theta = pt.atan2(xyzs[..., 2], xyzs[..., 0])
    length = pt.linalg.norm(xyzs, axis=-1)
    phi = pt.acos(xyzs[..., 1] / length)
    return pt.stack([theta, phi], dim=-1)

  def cube_xys_to_equally_spaced_theta_phi(xys): # shape: [..., 2], range: [-1, 1]
    theta = -xys[...,0] * math.pi / 4
    l = (pt.tan(theta) ** 2 + 1) ** 0.5
    phi = pt.atan2(l, xys[...,1])
    return pt.stack([theta, phi], dim=-1)

  num = resolution
  step = 2 / num
  offset = step / 2
  xs = pt.linspace(-1 + offset, 1 - offset, num)
  ys, xs = pt.meshgrid(xs, xs)
  xys = pt.dstack([xs, ys])
  theta_phis = cube_xys_to_equally_spaced_theta_phi(xys)

  XS = pt.sin(theta_phis[...,1]) * pt.cos(theta_phis[...,0])
  YS = pt.cos(theta_phis[...,1])
  ZS = pt.sin(theta_phis[...,1]) * pt.sin(theta_phis[...,0])
  pixel_dirs_x = pt.dstack([XS, YS, ZS])
  pixel_dirs_x_neg = pt.dstack([-XS, YS, -ZS])
  pixel_dirs_y = pt.dstack([-ZS, -XS, YS])
  pixel_dirs_y_neg = pt.dstack([-ZS, XS, -YS])
  pixel_dirs_z = pt.dstack([-ZS, YS, XS])
  pixel_dirs_z_neg = pt.dstack([ZS, YS, -XS])
  pixel_theta_phis = pt.stack([
    xyzs_to_theta_phi(pixel_dirs_x),
    xyzs_to_theta_phi(pixel_dirs_x_neg),
    xyzs_to_theta_phi(pixel_dirs_y),
    xyzs_to_theta_phi(pixel_dirs_y_neg),
    xyzs_to_theta_phi(pixel_dirs_z),
    xyzs_to_theta_phi(pixel_dirs_z_neg)], axis=0).cuda()

  basis_datas = []
  for postfix, theta_phi_ele in zip(
    ['x', 'x_neg', 'y', 'y_neg', 'z', 'z_neg'], pixel_theta_phis):
    pos_freq_viewing = model_specular.pos_freq_viewing
    while len(pos_freq_viewing.shape) != len(theta_phi_ele.shape) + 1:
      pos_freq_viewing = pos_freq_viewing.squeeze(0)
    theta_phi_ele[...,0] /= math.pi
    theta_phi_ele[...,1] = theta_phi_ele[...,1] / (math.pi * 0.5) - 1
    dxy = theta_phi_ele[..., None] * pos_freq_viewing
    dxy = dxy.view(list(dxy.shape[:-2]) + [2 * basis_order])
    dxy = pt.cat([pt.sin(dxy), pt.cos(dxy)], -1)
    basis = pt.tanh(model_specular.seq_basis(dxy))
    basis = basis.permute([2, 0, 1]).to(pt.float16)
    pt.cuda.empty_cache()
    basis_datas.append(basis.cpu().numpy())
    if DEBUG:
      for i in range(basis_out):
        im = ((basis[i] * 0.5 + 0.5) * 255 + 0.5).to(pt.uint8).cpu().numpy()
        cv2.imwrite(str(output_folder / f'basis_{postfix}_{i:02d}.png'), im)

  basis_datas = np.array(basis_datas)
  for i in range(basis_out):
    np.save(output_folder / f'basis_{i:02d}', basis_datas[:,i,...])

# export nex k(range: [-1, 1]) and alpha(range: [0, 1])
def export_nek_k(model, output_folder, msi_w, msi_h, dmin, dmax, msi_sublayers):
  phis, thetas = pt.meshgrid(
    pt.linspace(-1, 1, msi_h),
    pt.linspace(-1, 1, msi_w))
  theta_phi_coords = pt.dstack([thetas, phis]).cuda()
  pos_freq = model.pos_freq
  while len(pos_freq.shape) != len(theta_phi_coords.shape) + 1:
    pos_freq = pos_freq.squeeze(0)
  vxy = theta_phi_coords[..., None] * pos_freq
  vxy = vxy.view(list(vxy.shape[:-2]) + [2 * pos_level])
  alphas = []
  basis_k_datas = []
  for layer_i, radius in enumerate(model.radius):
    print(f'processing layer {layer_i}', end='\r')
    depth_freq = model.depth_freq
    while len(depth_freq.shape) != len(radius.shape) + 1:
      depth_freq = depth_freq.squeeze(0)
    vz = ((radius - dmin) / (dmax - dmin) * 2 - 1) * depth_freq
    vz = vz.expand(list(vxy.shape[:-1]) + [-1])
    vxyz = pt.cat([vxy, vz], -1)
    vxyz_coords = pt.cat([pt.sin(vxyz), pt.cos(vxyz)], -1)
    k = []
    row_is = pt.arange(0, vxyz_coords.shape[0], 200).tolist()
    row_is.append(vxyz_coords.shape[0])
    for i in range(len(row_is) - 1):
      k.append(model.seq_k(vxyz_coords[row_is[i]:row_is[i+1]]))
      pt.cuda.empty_cache()
    k = pt.cat(k, dim=0)
    alphas.append(pt.sigmoid(k[..., 0] - sigmoid_offset).cpu().numpy())
    if DEBUG:
      im = (alphas[-1] * 255 + 0.5).astype(np.uint8)
      cv2.imwrite(str(output_folder / f'msi_alpha_{layer_i:03d}.png'), im)
    if layer_i % msi_sublayers == 0:
      basis_k = pt.tanh(k[..., 1:])
      basis_k = basis_k.view(list(basis_k.shape[:-1]) + [3,-1])
      basis_k = basis_k.permute([3, 0, 1, 2]).to(pt.float16)
      basis_k_datas.append(basis_k.cpu().numpy())
      if DEBUG:
        im_all = ((basis_k * 0.5 + 0.5) * 255 + 0.5).to(pt.uint8).cpu().numpy()
        for i in range(basis_out):
          im = cv2.cvtColor(im_all[i], cv2.COLOR_RGB2BGR)
          cv2.imwrite(str(output_folder / f'basis_k_{layer_i:03d}_{i:02d}.png'), im)

  np.save(output_folder / 'msi_alpha', np.array(alphas))

  basis_k_datas = np.array(basis_k_datas)
  for i in range(basis_out):
    np.save(output_folder / f'basis_k_{i:02d}', basis_k_datas[:,i,...])

def load_model(filepath, dmin, dmax,
  w, h, theta_half_degree_range, phi_hald_degree_range):
  model = NexMsi(dmin, dmax, w, h,
    math.radians(theta_half_degree_range), math.radians(phi_hald_degree_range))

  checkpoint = pt.load(filepath)
  model.load_state_dict(checkpoint['model_state_dict'])
  epoch = checkpoint['epoch']
  print(f'Loading checkpoint @ Epoch {epoch}')
  return model

def run(config_file, model_ckpt_file, output_folder):
  with open(config_file) as f:
    config_data = json.load(f)
  dmin = config_data['dmin']
  dmax = config_data['dmax']
  msi_resolution_w = config_data['msi_resolution_w']
  msi_resolution_h = config_data['msi_resolution_h']
  theta_half_degree_range = config_data['theta_half_degree_range']
  phi_half_degree_range = config_data['phi_half_degree_range']

  model = load_model(model_ckpt_file,
    dmin, dmax, msi_resolution_w, msi_resolution_h,
    theta_half_degree_range, phi_half_degree_range)
  model.eval()
  pt.cuda.empty_cache()

  output_folder = Path(output_folder)
  output_folder.mkdir(parents=True, exist_ok=True)

  params_data = {}
  params_data['msi_layers'] = len(model.msi_rgb)
  params_data['total_layers'] = len(model.radius)
  params_data['radii'] = model.radius.squeeze().tolist()
  params_data['theta_half_degree_range'] = theta_half_degree_range
  params_data['phi_half_degree_range'] = phi_half_degree_range
  params_data['msi_h'] = model.msi_rgb.shape[-3]
  params_data['msi_w'] = model.msi_rgb.shape[-2]
  params_data['basis_out'] = basis_out
  with open(output_folder / 'params.json', 'w') as f:
    json.dump(params_data, f)

  cubemap_resolution = int(model.msi_rgb.shape[-2] * 45 / theta_half_degree_range + 0.5)
  with pt.no_grad():
    export_msi_images(model.msi_rgb, output_folder)
    del model.msi_rgb
    pt.cuda.empty_cache()
    export_nex_bases(model.specular, output_folder, cubemap_resolution)
    del model.specular
    pt.cuda.empty_cache()
    msi_sublayers = params_data['total_layers'] // params_data['msi_layers']
    export_nek_k(model, output_folder,
      params_data['msi_w'], params_data['msi_h'], dmin, dmax, msi_sublayers)
    pt.cuda.empty_cache()

if __name__ == '__main__':
  if len(sys.argv) < 4:
    print('export_nex_data config_file model_ckpt output_folder')
    exit(-1)

  run(sys.argv[1], sys.argv[2], sys.argv[3])
